{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc47a932-9b5e-4b24-84f1-68813e1a8fff",
   "metadata": {
    "is_executing": true
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import glob\n",
    "import os\n",
    "\n",
    "config = dict()\n",
    "\n",
    "model_config = dict()\n",
    "model_config[\"name\"] = \"DynUNet\"  # network model name from MONAI\n",
    "# set the network hyper-parameters\n",
    "model_config[\"in_channels\"] = 4  # 4 input images for the BraTS challenge\n",
    "model_config[\"out_channels\"] = 3   # whole tumor, tumor core, enhancing tumor\n",
    "model_config[\"spatial_dims\"] = 3   # 3D input images\n",
    "model_config[\"deep_supervision\"] = False  # do not check outputs of lower layers\n",
    "model_config[\"strides\"] = [[1, 1, 1], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2]][:-1]  # number of downsampling convolutions\n",
    "model_config[\"filters\"] = [64, 96, 128, 192, 256, 384, 512, 768, 1024][:len(model_config[\"strides\"])]  # number of filters per layer\n",
    "model_config[\"kernel_size\"] = [[3, 3, 3]] * len(model_config[\"strides\"])  # size of the convolution kernels per layer\n",
    "model_config[\"upsample_kernel_size\"] = model_config[\"strides\"][1:]  # should be the same as the strides\n",
    "\n",
    "# put the model config in the main config\n",
    "config[\"model\"] = model_config\n",
    "\n",
    "config[\"optimizer\"] = {'name': 'Adam', \n",
    "                       'lr': 0.001}  # initial learning rate\n",
    "\n",
    "# define the loss\n",
    "config[\"loss\"] = {'name': 'DiceLoss', # from Monai\n",
    "                  'include_background': True,  # we do not have a label for the background, so this should be true (by \"include background\" monai means include channel 0)\n",
    "                  'sigmoid': True}  # transform the model logits to activations\n",
    "\n",
    "# set the cross validation parameters\n",
    "config[\"cross_validation\"] = {'n_folds': 5,  # number of cross validation folds\n",
    "                              'seed': 25},  # seed to make the generation of cross validation folds consistent across different trials\n",
    "# set the scheduler parameters\n",
    "config[\"scheduler\"] = {'name': 'ReduceLROnPlateau', \n",
    "                       'patience': 10,  # wait 10 epochs with no improvement before reducing the learning rate\n",
    "                       'factor': 0.5,   # multiply the learning rate by 0.5\n",
    "                       'min_lr': 1e-08}  # stop reducing the learning rate once it gets to 1e-8\n",
    "\n",
    "# set the dataset parameters\n",
    "config[\"dataset\"] = {'name': 'SegmentationDatasetPersistent',  # 'Persistent' means that it will save the preprocessed outputs generated during the first epoch\n",
    "# However, using 'Persistent', does also increase the time of the first epoch compared to the other epochs, which should run faster\n",
    "  'desired_shape': [128, 128, 128],  # resize the images to this shape, increase this to get higher resolution images (increases computation time and memory usage)\n",
    "  'labels': [2, 1, 4],  # 1: necrotic center; 2: edema, 3: enhancing tumor\n",
    "  'setup_label_hierarchy': True,  # changes the labels to whole tumor (2, 1, 4), tumor core (1, 4), and enhancing tumor (4) to be consistent with the challenge\n",
    "  'normalization': 'NormalizeIntensityD',  # z score normalize the input images to zero mean unit standard deviation\n",
    "  'normalization_kwargs': {'channel_wise': True, \"nonzero\": False},  # perform the normalization channel wise and include the background\n",
    "  'resample': True,  # resample the images when resizing them, otherwise the resize could crop out regions of interest\n",
    "  'crop_foreground': True,  # crop the foreground of the images\n",
    "                    }\n",
    "config[\"training\"] = {'batch_size': 1,  # number of image/label pairs to read at a time during training\n",
    "  'validation_batch_size': 1,  # number of image/label pairs to read at atime during validation\n",
    "  'amp': False,  # don't set this to true unless the model you are using is setup to use automatic mixed precision (AMP)\n",
    "  'early_stopping_patience': None,  # stop the model early if the validaiton loss stops improving\n",
    "  'n_epochs': 250,  # number of training epochs, reduce this if you don't want training to run as long\n",
    "  'save_every_n_epochs': None,  # save the model every n epochs (otherwise only the latest model will be saved)\n",
    "  'save_last_n_models': None,  # save the last n models \n",
    "  'save_best': True}  # save the model that has the best validation loss\n",
    "\n",
    "# get the training filenames\n",
    "config[\"training_filenames\"] = list()\n",
    "\n",
    "# if your BraTS data is stored somewhere else, change this code to fetch that data\n",
    "for subject_folder in sorted(glob.glob(\"BraTS2020_TrainingData/MICCAI_BraTS2020_TrainingData/*\")):\n",
    "    if not os.path.isdir(subject_folder):\n",
    "        continue\n",
    "    image_filenames = sorted(glob.glob(os.path.join(subject_folder, \"*.nii\")))\n",
    "    for i in range(len(image_filenames)):\n",
    "        if \"seg\" in image_filenames[i].lower():\n",
    "            label = image_filenames.pop(i)\n",
    "            break\n",
    "    assert len(image_filenames) == 4\n",
    "    config[\"training_filenames\"].append({\"image\": image_filenames, \"label\": label})\n",
    "\n",
    "\n",
    "config[\"bratsvalidation_filenames\"] = list()  # \"validation_filenames\" is reserved for the cross-validation, so we will call this bratsvalidation_filenames\n",
    "for subject_folder in sorted(glob.glob(\"BraTS2020_ValidationData/MICCAI_BraTS2020_ValidationData/*\")):\n",
    "    if not os.path.isdir(subject_folder):\n",
    "        continue\n",
    "    image_filenames = sorted(glob.glob(os.path.join(subject_folder, \"*.nii\")))\n",
    "    assert len(image_filenames) == 4\n",
    "    config[\"bratsvalidation_filenames\"].append({\"image\": image_filenames})\n",
    "\n",
    "\n",
    "with open(\"./brats2020_config.json\", \"w\") as op:\n",
    "    json.dump(config, op, indent=4)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (brats)",
   "language": "python",
   "name": "brats"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
